import numpy as np
import torch.utils.data as data
from glob import glob
import cv2
import scipy.io as sio
import torch
from utils.presets import SimpleTransform
import yaml
import copy


class Air2land(data.Dataset):
    EVAL_JOINTS = list(range(6))
    joint_pairs = [[2, 0], [1, 2], [2, 3], [2, 4], [2, 5]]

    def __init__(self, cfg, mode="train", dpg=False):
        if mode == "train":
            self.config = cfg["DATASET"]["TRAIN"]
        else:
            self.config = cfg["DATASET"]["VAL"]
        self.root = self.config["ROOT"]
        if mode == "train":
            self.scene = self.root + "train.txt"
        elif mode == "test":
            self.scene = self.root + "test.txt"
        with open(self.scene, "r") as f:
            self.scene_list = f.readlines()
            self.scene_list = [x.replace("\n", "") for x in self.scene_list]
        self.imgs = []
        # self.anns = []
        print("{}集合中包含场景：{}个".format(mode, len(self.scene_list)))
        for scene in self.scene_list:
            img_path = glob(scene + "/*/*/" + "*-color.jpg")
            self.imgs.extend(img_path)
        print("{}集合中包含图片：{}个".format(mode, str(len(self.imgs))))

        if 'AUG' in self.config.keys():
            self._scale_factor = self.config['AUG']['SCALE_FACTOR']
            self._rot = self.config['AUG']['ROT_FACTOR']
        else:
            self._scale_factor = 0
            self._rot = 0

        self._input_size = cfg["DATA_PRESET"]['IMAGE_SIZE']
        self._output_size = cfg["DATA_PRESET"]['HEATMAP_SIZE']
        self._sigma = cfg["DATA_PRESET"]['SIGMA']
        self._dpg = dpg
        self._train = True if mode == "train" else False
        self._loss_type = cfg.get('LOSS_TYPE', 'MSELoss')

        self.transformation = SimpleTransform(
            self, scale_factor=self._scale_factor,
            input_size=self._input_size,
            output_size=self._output_size,
            rot=self._rot, sigma=self._sigma,
            train=self._train, add_dpg=self._dpg,
            loss_type=self._loss_type
        )

    def load_mat(self, ann_path):
        data = sio.loadmat(ann_path)
        PM = np.array(data["relative_poses"])
        pts2d = np.array(data["key_points"])
        bbox = np.array(data["bounding_boxes"])[0]
        x_min = bbox[0]
        y_min = bbox[1]
        w = bbox[2]
        h = bbox[3]
        x_max = x_min + w
        y_max = y_min + h
        bbox = np.array([x_min, y_min, x_max, y_max])
        label = {"pts2d": pts2d, "bbox": bbox, "relative_pose": PM}
        return label

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, idx):
        img = cv2.imread(self.imgs[idx])
        img_ori = copy.deepcopy(img)
        ann_path = self.imgs[idx].replace("color", "meta").replace("jpg", "mat")
        label = self.load_mat(ann_path)
        kp = copy.deepcopy(label["pts2d"])
        PM = copy.deepcopy(label["relative_pose"])
        img, label, label_mask, bbox = self.transformation(img, label)
        return img_ori, img, label, label_mask, idx, bbox, kp, PM, self.imgs[idx]

# if __name__ =="__main__":
#     with open("/media/liyuke/share/AAA/AlphaPose-master/configs/coco/resnet/air2land_256x192_res50_lr1e-3_2x.yaml", 'r') as f:
#         config = yaml.load(f, Loader=yaml.FullLoader)
#     dataset = Air2land(config,"train")
#     out = dataset.__getitem__(1)
